import queue
import random

from concurrency.taskBySelf import TaskBySelf
from concurrency.workerThread import WorkerThread


class ThreadPool:
    """线程池类，管理多个工作线程"""
    def __init__(self, num_threads):
        self.task_queue = queue.Queue()  # 任务队列
        self.threads = [WorkerThread(self.task_queue) for _ in range(num_threads)]
        self.tasks = {}  # 用于存储任务的唯一标识符和任务对象

    def start(self):
        """启动线程池"""
        for thread in self.threads:
            thread.start()
        print(f"线程池启动，包含 {len(self.threads)} 个工作线程")

    def add_task(self, func, *args, **kwargs):
        """向任务队列中添加任务，支持传参"""
        task_id = random.randint(1000, 9999)  # 生成一个随机任务ID
        task = TaskBySelf(func, args, kwargs)
        self.tasks[task_id] = task  # 将任务存储到任务字典中
        self.task_queue.put(task)
        print(f"任务 {task_id} 添加到队列: {func.__name__}，参数: {args}, {kwargs}")
        return task_id

    def stop_task(self, task_id):
        """停止指定任务"""
        if task_id in self.tasks:
            task = self.tasks[task_id]
            task.stop()
            print(f"任务 {task_id} 已停止")
        else:
            print(f"任务 {task_id} 不存在")

    def stop_all(self):
        """停止所有线程和任务"""
        for task in self.tasks.values():
            task.stop()
        for thread in self.threads:
            thread.join()  # 等待所有线程停止
        print("线程池已停止")

    def wait_completion(self):
        """等待所有任务完成"""
        self.task_queue.join()
        print("所有任务已完成")